#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
SepMark 批量水印嵌入脚本
专门用于给数据集批量嵌入水印
"""

import os
import torch
import numpy as np
from PIL import Image
import yaml
from easydict import EasyDict
from torchvision import transforms
from network.Dual_Mark import Network
import argparse
import time
from tqdm import tqdm


class WatermarkEmbedder:
    def __init__(self, config_path=r"D:\deepfake\SepMark-main\SepMark-main\cfg\test_DualMark.yaml"):
        """
        初始化水印嵌入器

        Args:
            config_path: 配置文件路径
        """
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"使用设备: {self.device}")

        # 加载配置
        self.load_config(config_path)

        # 初始化网络
        self.setup_network()

        # 设置图像转换
        self.transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

        # 反向转换（用于保存图像）
        self.inverse_transform = transforms.Compose([
            transforms.Normalize([-1, -1, -1], [2, 2, 2]),  # 反归一化: [-1,1] -> [0,1]
            transforms.ToPILImage()
        ])

        print("水印嵌入器初始化完成")

    def load_config(self, config_path):
        """加载配置文件"""
        with open(config_path, 'r') as f:
            test_args = EasyDict(yaml.load(f, Loader=yaml.SafeLoader))

        self.result_folder = "results/" + test_args.result_folder
        self.model_epoch = test_args.model_epoch
        self.strength_factor = test_args.strength_factor

        # 加载训练配置
        with open(self.result_folder + '/train_DualMark.yaml', 'r') as f:
            train_args = EasyDict(yaml.load(f, Loader=yaml.SafeLoader))

        self.message_length = train_args.message_length
        self.message_range = train_args.message_range
        self.attention_encoder = train_args.attention_encoder
        self.attention_decoder = train_args.attention_decoder
        self.weight = train_args.weight
        self.noise_layers_R = train_args.noise_layers.pool_R
        self.noise_layers_F = train_args.noise_layers.pool_F

        print(f"配置加载完成: 模型epoch={self.model_epoch}, 强度因子={self.strength_factor}")

    def setup_network(self):
        """设置网络结构并加载模型"""
        # 初始化网络
        self.network = Network(
            self.message_length,
            self.noise_layers_R,
            self.noise_layers_F,
            self.device,
            1,  # batch_size = 1
            0.0002,  # lr (推理时不用)
            0.5,  # beta1 (推理时不用)
            self.attention_encoder,
            self.attention_decoder,
            self.weight
        )

        # 加载预训练模型
        model_path = self.result_folder + f"models/EC_{self.model_epoch}.pth"
        if os.path.exists(model_path):
            self.network.load_model_ed(model_path)
            print(f"模型加载成功: {model_path}")
        else:
            raise FileNotFoundError(f"模型文件未找到: {model_path}")

        # 设置为评估模式
        self.network.encoder_decoder.eval()

    def load_image(self, image_path):
        """
        加载并预处理图像

        Args:
            image_path: 图像路径

        Returns:
            处理后的图像tensor
        """
        image = Image.open(image_path).convert('RGB')
        image_tensor = self.transform(image).unsqueeze(0).to(self.device)
        return image_tensor

    def generate_watermark_message(self, seed=None):
        """
        生成水印消息

        Args:
            seed: 随机种子（可选）

        Returns:
            水印消息tensor
        """
        if seed is not None:
            np.random.seed(seed)

        message = torch.Tensor(
            np.random.choice([-self.message_range, self.message_range],
                             (1, self.message_length))
        ).to(self.device)

        return message

    def embed_watermark(self, image, message):
        """
        在图像中嵌入水印

        Args:
            image: 输入图像tensor
            message: 水印消息tensor

        Returns:
            嵌入水印后的图像tensor
        """
        with torch.no_grad():
            # 使用编码器嵌入水印
            encoded_image = self.network.encoder_decoder.module.encoder(image, message)

            # 应用强度因子
            watermarked_image = image + (encoded_image - image) * self.strength_factor

            return watermarked_image

    def save_image(self, image_tensor, save_path):
        """
        保存图像

        Args:
            image_tensor: 图像tensor
            save_path: 保存路径
        """
        # 确保目录存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        # 转换并保存
        image_pil = self.inverse_transform(image_tensor.squeeze(0).cpu())
        image_pil.save(save_path)

    def embed_dataset(self, input_dir, output_dir, watermark_seed=42, save_original=True):
        """
        批量给数据集嵌入水印

        Args:
            input_dir: 输入图像目录
            output_dir: 输出目录
            watermark_seed: 水印种子（所有图像使用相同水印）
            save_original: 是否同时保存原图
        """
        # 支持的图像格式
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}

        # 获取所有图像文件
        image_files = []
        for root, dirs, files in os.walk(input_dir):
            for file in files:
                if any(file.lower().endswith(ext) for ext in image_extensions):
                    rel_path = os.path.relpath(os.path.join(root, file), input_dir)
                    image_files.append(rel_path)

        print(f"找到 {len(image_files)} 张图像")

        # 生成水印消息
        message = self.generate_watermark_message(watermark_seed)
        print(f"水印消息生成完成 (种子: {watermark_seed})")

        # 创建输出目录
        watermarked_dir = os.path.join(output_dir, 'watermarked')
        if save_original:
            original_dir = os.path.join(output_dir, 'original')

        # 处理每张图像
        successful_count = 0
        failed_count = 0

        for image_file in tqdm(image_files, desc="嵌入水印"):
            try:
                # 构建路径
                input_path = os.path.join(input_dir, image_file)
                watermarked_path = os.path.join(watermarked_dir, image_file)

                # 加载图像
                image = self.load_image(input_path)

                # 嵌入水印
                watermarked_image = self.embed_watermark(image, message)

                # 保存水印图像
                self.save_image(watermarked_image, watermarked_path)

                # 保存原图（如果需要）
                if save_original:
                    original_path = os.path.join(original_dir, image_file)
                    self.save_image(image, original_path)

                successful_count += 1

            except Exception as e:
                print(f"处理 {image_file} 时出错: {e}")
                failed_count += 1

        print(f"\n批量水印嵌入完成!")
        print(f"成功处理: {successful_count} 张")
        print(f"失败: {failed_count} 张")
        print(f"水印图像保存在: {watermarked_dir}")
        if save_original:
            print(f"原图保存在: {original_dir}")

        # 保存水印信息
        watermark_info = {
            'watermark_seed': watermark_seed,
            'message_length': self.message_length,
            'message_range': self.message_range,
            'strength_factor': self.strength_factor,
            'model_epoch': self.model_epoch,
            'total_images': len(image_files),
            'successful_count': successful_count,
            'failed_count': failed_count
        }

        import json
        with open(os.path.join(output_dir, 'watermark_info.json'), 'w') as f:
            json.dump(watermark_info, f, indent=2)

        print(f"水印信息保存在: {os.path.join(output_dir, 'watermark_info.json')}")


def main():
    """主函数"""

    # ===========================================
    # 直接在这里设置固定的输入和输出目录
    # ===========================================

    # 固定输入目录 - 修改为您的图片目录
    FIXED_INPUT_DIR = r"D:\deepfake\CelebAMask-HQ\CelebAMask-HQ\test"

    # 固定输出目录 - 您想保存水印图片的路径
    FIXED_OUTPUT_DIR = r"D:\deepfake\CelebAMask-HQ\CelebAMask-HQ\marked"

    # 其他固定参数
    WATERMARK_SEED = 42  # 水印种子
    SAVE_ORIGINAL = True  # 是否同时保存原图
    CONFIG_PATH = 'cfg/test_DualMark.yaml'  # 配置文件路径

    # ===========================================

    print(f"输入目录: {FIXED_INPUT_DIR}")
    print(f"输出目录: {FIXED_OUTPUT_DIR}")

    # 检查输入目录
    if not os.path.exists(FIXED_INPUT_DIR):
        print(f"错误: 输入目录不存在: {FIXED_INPUT_DIR}")
        print("请创建目录并放入图片，或修改 FIXED_INPUT_DIR 路径")
        return

    # 创建水印嵌入器
    print("正在初始化水印嵌入器...")
    embedder = WatermarkEmbedder(CONFIG_PATH)

    # 执行批量嵌入
    print("开始批量水印嵌入...")
    embedder.embed_dataset(
        input_dir=FIXED_INPUT_DIR,
        output_dir=FIXED_OUTPUT_DIR,
        watermark_seed=WATERMARK_SEED,
        save_original=SAVE_ORIGINAL
    )

    print("水印嵌入完成！")

if __name__ == '__main__':
    main()